/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: 集合通信算子信息库
 * Author: jiangchen
 * Create: 2020-6-10
 */

#include "hvd_ops_kernel_info_store.h"
#include <securec.h>
#include <functional>
#include "external/graph/tensor.h"
#include "graph/utils/attr_utils.h"
#include "graph/utils/tensor_utils.h"
#include "graph/debug/ge_attr_define.h"
#include "framework/common/fmk_error_codes.h"
#include "external/graph/types.h"
#include "hccl/hcom.h"
#include "comm.h"
#include "hvd_adapter.h"
#include "runtime/kernel.h"

namespace hccl {
HvdOpsKernelInfoStore::HvdOpsKernelInfoStore() : eventMap_(), callbackVec_()
{
}

HvdOpsKernelInfoStore::~HvdOpsKernelInfoStore()
{
    for (u32 i = 0; i < callbackVec_.size(); i++) {
        if (callbackVec_[i] != nullptr) {
            delete callbackVec_[i];
            callbackVec_[i] = nullptr;
        }
    }
    callbackVec_.clear();
    HcclResult ret = DestroyEvent();
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[Store][KernelInfo]DestroyEvent failed in HvdOpsKernelInfoStore Destructor, ret[%d].", ret);
    }
}

HcclResult HvdOpsKernelInfoStore::SetCustomKernelInfo(ge::OpInfo &opinfo, std::map<string, ge::OpInfo> &infos) const
{
    opinfo.opKernelLib = HVD_OPS_LIB_NAME;
    for (u32 index = 0; index < HVD_SUPPORTED_OP_TYPE.size(); index++) {
        HCCL_INFO("op[%s]: engine[%s] opKernelLib[%s] computeCost[%d] flagPartial[%d] flagAsync[%d]", \
            HVD_SUPPORTED_OP_TYPE[index].c_str(), opinfo.engine.c_str(), opinfo.opKernelLib.c_str(), \
            opinfo.computeCost, opinfo.flagPartial, opinfo.flagAsync);
        // hvdcallbackAllreduce set atomic clean to keep the data clean
        if (HVD_SUPPORTED_OP_TYPE[index] == HVD_KERNEL_OP_TYPE_ALLREDUCE) {
            opinfo.isAtomic = true;
        } else {
            opinfo.isAtomic = false;
        }
        infos.insert(std::pair<string, ge::OpInfo>(HVD_SUPPORTED_OP_TYPE[index], opinfo));
    }
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::GetSupportedOP(std::vector<std::string> &hcclSupportOp) const
{
    hcclSupportOp.assign(HVD_SUPPORTED_OP_TYPE.begin(), HVD_SUPPORTED_OP_TYPE.end());
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::HCCLOpsKernel(const ge::GETaskInfo &task, const std::string &sCollectiveType)
{
    if (sCollectiveType == HVD_KERNEL_OP_TYPE_BROADCAST) {
        CHK_RET(HvdBroadcastOpKernel(task, sCollectiveType));
    } else if (sCollectiveType == HVD_KERNEL_OP_TYPE_ALLREDUCE) {
        CHK_RET(HvdAllReduceOpKernel(task, sCollectiveType));
    } else if (sCollectiveType == HVD_KERNEL_OP_TYPE_ALLGATHER) {
        CHK_RET(HvdAllGatherOpKernel(task, sCollectiveType));
    } else if (sCollectiveType == HVD_KERNEL_OP_TYPE_WAIT) {
        CHK_RET(HvdWaitOpKernel(task, sCollectiveType));
    } else {
        HCCL_ERROR("[HCCLOps][Kernel]errNo[0x%016llx]Op type[%s] is invalid.", HCOM_ERROR_CODE(HCCL_E_PARA),
            sCollectiveType.c_str());
        return HCCL_E_PARA;
    }
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::HvdBroadcastOpKernel(const ge::GETaskInfo &task, const string &sCollectiveType)
{
    HcclDataType dataType;
    uintptr_t inputAddr = 0;
    u64 count;
    u32 root;
    u32 eventID;
    rtEvent_t event = nullptr;
    rtStream_t streamMain;
    std::vector<ge::GETaskKernelHcclInfo> hcclInfos = task.kernelHcclInfo;
    auto hvdCallBackInfo = new (std::nothrow) HvdCallBackInfo;
    CHK_PTR_NULL(hvdCallBackInfo);
    callbackVec_.push_back(hvdCallBackInfo);
    hvdCallBackInfo->vectorSize = task.kernelHcclInfo.size();
    hvdCallBackInfo->hcclType = sCollectiveType;

    CHK_RET(GetRootFromTaskInfo(hcclInfos[0], root));  // pick first to get root rank

    hvdCallBackInfo->rootId = root;
        CHK_RET(GetEventIDFromTaskInfo(task, eventID));  // same eventID reveals the pair of hvdCallback op and wait op

    CHK_RET(GetEvent(eventID, event));
    hvdCallBackInfo->event = event;

    hvdCallBackInfo->kernelHvdInfo.resize(task.kernelHcclInfo.size());
    for (u32 i = 0; i < hcclInfos.size(); i++) {
        CHK_RET(GetInputAddrFromTaskInfo(hcclInfos[i], inputAddr));
        CHK_RET(GetDataTypeFromTaskInfo(task, dataType));
        CHK_RET(GetCountFromTaskInfo(hcclInfos[i], count));

        hvdCallBackInfo->kernelHvdInfo[i].name = hcclInfos[i].input_name;
        hvdCallBackInfo->kernelHvdInfo[i].inputDataAddr = reinterpret_cast<void *>(inputAddr);
        hvdCallBackInfo->kernelHvdInfo[i].count = count;
        hvdCallBackInfo->kernelHvdInfo[i].dataType = dataType;
        hvdCallBackInfo->kernelHvdInfo[i].dims = hcclInfos[i].dims;
    }
    CHK_RET(GetStreamMainFromTaskInfo(task, streamMain));

    s32 deviceID;
    rtError_t rtRet = rtGetDevice(&deviceID);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[Op][Kernel]rtGetDevice failed:[%d]", rtRet), HCCL_E_RUNTIME);
    hvdCallBackInfo->deviceID = deviceID;
    rtRet = rtSetDevice(deviceID);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[Op][Kernel]rtSetDevice failed:[%d]", rtRet), HCCL_E_RUNTIME);

    CHK_RET(g_hvdAdapterGlobal.HvdAdapterInit(streamMain, deviceID));  // device should be set in new thread

    // Launch rtCallback task which will execute CallbackFunc later
    rtRet = rtCallbackLaunch(HvdOpCallbackFunc, hvdCallBackInfo, streamMain, true);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[Op][Kernel]rtCallbackLaunch failed:[%d]", rtRet),
        HCCL_E_RUNTIME);
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::HvdAllReduceOpKernel(const ge::GETaskInfo &task, const string &sCollectiveType)
{
    HcclDataType dataType;
    uintptr_t inputAddr = 0;
    uintptr_t outputAddr = 0;
    u64 count;
    HcclReduceOp reduceType;
    u32 eventID;
    rtEvent_t event = nullptr;
    rtStream_t streamMain;
    std::vector<ge::GETaskKernelHcclInfo> hcclInfos = task.kernelHcclInfo;
    auto hvdCallBackInfo = new (std::nothrow) HvdCallBackInfo;
    CHK_PTR_NULL(hvdCallBackInfo);
    callbackVec_.push_back(hvdCallBackInfo);
    hvdCallBackInfo->vectorSize = task.kernelHcclInfo.size();
    hvdCallBackInfo->hcclType = sCollectiveType;
    CHK_RET(GetReduceTypeFromTaskInfo(hcclInfos[0], reduceType)); // pick first to get reducetype;
    hvdCallBackInfo->opType = reduceType;

    CHK_RET(GetEventIDFromTaskInfo(task, eventID));
    CHK_RET(GetEvent(eventID, event));
    hvdCallBackInfo->event = event;

    hvdCallBackInfo->kernelHvdInfo.resize(task.kernelHcclInfo.size());
    for (u32 i = 0; i < hcclInfos.size(); i++) {
        CHK_RET(GetInputAddrFromTaskInfo(hcclInfos[i], inputAddr));
        CHK_RET(GetOutputAddrFromTaskInfo(hcclInfos[i], outputAddr));
        CHK_RET(GetDataTypeFromTaskInfo(task, dataType));
        CHK_RET(GetCountFromTaskInfo(hcclInfos[i], count));

        hvdCallBackInfo->kernelHvdInfo[i].name = hcclInfos[i].input_name;
        hvdCallBackInfo->kernelHvdInfo[i].inputDataAddr = reinterpret_cast<void *>(inputAddr);
        hvdCallBackInfo->kernelHvdInfo[i].outputDataAddr =  reinterpret_cast<void *>(outputAddr);
        hvdCallBackInfo->kernelHvdInfo[i].count = count;
        hvdCallBackInfo->kernelHvdInfo[i].dataType = dataType;
        hvdCallBackInfo->kernelHvdInfo[i].dims = hcclInfos[i].dims;
    }
    CHK_RET(GetStreamMainFromTaskInfo(task, streamMain));

    s32 deviceID;
    rtError_t rtRet = rtGetDevice(&deviceID);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[AllReduceOp][Kernel]rtGetDevice failed:[%d].", rtRet),
        HCCL_E_RUNTIME);
    hvdCallBackInfo->deviceID = deviceID;
    rtRet = rtSetDevice(deviceID);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[AllReduceOp][Kernel]rtSetDevice failed:[%d]", rtRet),
        HCCL_E_RUNTIME);
    CHK_RET(g_hvdAdapterGlobal.HvdAdapterInit(streamMain, deviceID));  // device should be set in new thread

    // Launch rtCallback task which will execute CallbackFunc later
    rtRet = rtCallbackLaunch(HvdOpCallbackFunc, hvdCallBackInfo, streamMain, false);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[AllReduceOp][Kernel]rtCallbackLaunch failed:[%d].", rtRet),
        HCCL_E_RUNTIME);
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::HvdAllGatherOpKernel(const ge::GETaskInfo &task, const string &sCollectiveType)
{
    HcclDataType dataType;
    uintptr_t inputAddr = 0;
    uintptr_t outputAddr = 0;
    u64 count;
    u32 eventID;
    rtEvent_t event = nullptr;
    rtStream_t streamMain;
    std::vector<ge::GETaskKernelHcclInfo> hcclInfos = task.kernelHcclInfo;
    auto hvdCallBackInfo = new (std::nothrow) HvdCallBackInfo;
    CHK_PTR_NULL(hvdCallBackInfo);
    callbackVec_.push_back(hvdCallBackInfo);
    hvdCallBackInfo->vectorSize = task.kernelHcclInfo.size();
    hvdCallBackInfo->hcclType = sCollectiveType;

    CHK_RET(GetEventIDFromTaskInfo(task, eventID));
    CHK_RET(GetEvent(eventID, event));
    hvdCallBackInfo->event = event;

    hvdCallBackInfo->kernelHvdInfo.resize(task.kernelHcclInfo.size());
    for (u32 i = 0; i < hcclInfos.size(); i++) {
        CHK_RET(GetInputAddrFromTaskInfo(hcclInfos[i], inputAddr));
        CHK_RET(GetOutputAddrFromTaskInfo(hcclInfos[i], outputAddr));
        CHK_RET(GetDataTypeFromTaskInfo(task, dataType));
        CHK_RET(GetCountFromTaskInfo(hcclInfos[i], count));

        hvdCallBackInfo->kernelHvdInfo[i].name = hcclInfos[i].input_name;
        hvdCallBackInfo->kernelHvdInfo[i].inputDataAddr = reinterpret_cast<void *>(inputAddr);
        hvdCallBackInfo->kernelHvdInfo[i].outputDataAddr = reinterpret_cast<void *>(outputAddr);
        hvdCallBackInfo->kernelHvdInfo[i].count = count;
        hvdCallBackInfo->kernelHvdInfo[i].dataType = dataType;
        hvdCallBackInfo->kernelHvdInfo[i].dims = hcclInfos[i].dims;
    }
    CHK_RET(GetStreamMainFromTaskInfo(task, streamMain));

    s32 deviceID;
    rtError_t rtRet = rtGetDevice(&deviceID);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[AllGatherOp][Kernel]rtGetDevice failed:[%d].", rtRet),
        HCCL_E_RUNTIME);
    hvdCallBackInfo->deviceID = deviceID;
    rtRet = rtSetDevice(deviceID);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[AllGatherOp][Kernel]rtSetDevice failed:[%d]", rtRet),
        HCCL_E_RUNTIME);

    CHK_RET(g_hvdAdapterGlobal.HvdAdapterInit(streamMain, deviceID));  // device should be set in new thread

    // Launch rtCallback task which will execute CallbackFunc later
    rtRet = rtCallbackLaunch(HvdOpCallbackFunc, hvdCallBackInfo, streamMain, false);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[AllGatherOp][Kernel]rtCallbackLaunch failed:[%d].",
        rtRet), HCCL_E_RUNTIME);
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::HvdWaitOpKernel(const ge::GETaskInfo &task, const string &sCollectiveType)
{
    u32 eventID;
    rtEvent_t event = nullptr;
    rtStream_t streamMain;

    CHK_RET(GetEventIDFromTaskInfo(task, eventID));
    CHK_RET(GetEvent(eventID, event));
    CHK_RET(GetStreamMainFromTaskInfo(task, streamMain));
    HCCL_DEBUG("tasktype:[%s]", sCollectiveType.c_str());

    s32 deviceID;
    rtError_t rtRet = rtGetDevice(&deviceID);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[WaitOp][Kernel]rtGetDevice failed:[%d].", rtRet),
        HCCL_E_RUNTIME);
    rtRet = rtSetDevice(deviceID);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[WaitOp][Kernel]rtSetDevice failed:[%d]", rtRet),
        HCCL_E_RUNTIME);
    // Launch event wait task which will block till the corresponding event record
    rtRet = rtStreamWaitEvent(streamMain, event);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[WaitOp][Kernel]rtStreamWaitEvent failed:[%d]", rtRet),
        HCCL_E_RUNTIME);
    // Once event wait success, reset it for next iteration
    rtRet = rtEventReset(event, streamMain);
    CHK_PRT_RET((rtRet != RT_ERROR_NONE), HCCL_ERROR("[WaitOp][Kernel]rtEventReset failed:[%d]", rtRet),
        HCCL_E_RUNTIME);

    return HCCL_SUCCESS;
}

ge::Status HvdOpsKernelInfoStore::LoadTask(ge::GETaskInfo &task)
{
    std::string sCollectiveType;
    // std::string sTag;
    HCCL_INFO("LoadTask Start.");
    g_hvdMode = true;
    CHK_PRT_RET(!(task.kernelHcclInfo.size() > 0), HCCL_ERROR("[Load][Task]HcclInfo in task is less than 1:[%d]",
        task.kernelHcclInfo.size()), ge::INTERNAL_ERROR);
    ge::GETaskKernelHcclInfo hcclInfo = task.kernelHcclInfo[0]; // get the task type from first item

    CHK_PRT_RET((task.type != RT_MODEL_TASK_HCCL), HCCL_ERROR("[Load][Task]errNo[0x%016llx] TaskType[%u] from"
        "taskinfo is invalid", HCOM_ERROR_CODE(HCCL_E_PARA), task.type), ge::INTERNAL_ERROR);

    HcclResult ret = GetCollectiveTypeFromTaskInfo(hcclInfo, sCollectiveType);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Load][Task]errNo[0x%016llx] load task failed. (get collective type fail)",
            HCOM_ERROR_CODE(ret)), ge::INTERNAL_ERROR);
    // exectue hvdCallbackKernel
    ret = HCCLOpsKernel(task, sCollectiveType);
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[Load][Task]errNo[0x%016llx] load task failed. op[%s]", HCOM_ERROR_CODE(ret),
            sCollectiveType.c_str());
        if (callbackVec_[callbackVec_.size() - 1] != nullptr) {
            delete callbackVec_[callbackVec_.size() - 1];
            callbackVec_[callbackVec_.size() - 1] = nullptr;
            callbackVec_.pop_back();
        }
        return ge::INTERNAL_ERROR;
    }
    HCCL_INFO("LoadTask success. Type is [%s]", sCollectiveType.c_str());
    return ge::SUCCESS;
}

ge::Status HvdOpsKernelInfoStore::UnloadTask(ge::GETaskInfo &task)
{
    return ge::SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::GetEventIDFromTaskInfo(const ge::GETaskInfo &task, u32 &eventID)
{
    HvdKernelInfoPrivateDef *privateDefBuf = reinterpret_cast<HvdKernelInfoPrivateDef *>(task.privateDef);
    eventID = privateDefBuf->eventID;
    HCCL_DEBUG("get src eventID[%u] from task info success.", eventID);
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::GetEvent(const u32 &eventID, rtEvent_t &event)
{
    rtError_t ret;
    rtEvent_t newEvent;
    auto iter = eventMap_.find(eventID);
    HCCL_DEBUG("eventMap_ size:[%d]", eventMap_.size());
    if (iter == eventMap_.end()) {
        ret = rtEventCreateWithFlag(&newEvent, 1);
        if (ret != RT_ERROR_NONE) {
            HCCL_ERROR("[Get][Event]create event failed.");
            return HCCL_E_RUNTIME;
        }
        eventMap_.insert(std::make_pair(eventID, newEvent));
        HCCL_DEBUG("create an event id [%u] and insert to eventMap_", eventID);
    } else {
        newEvent = iter->second;
        HCCL_DEBUG("get an event id [%u] from eventMap_ size", eventID);
    }
    event = newEvent;
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::DestroyEvent()
{
    rtError_t ret;
    for (auto each : eventMap_) {
        ret = rtEventDestroy(each.second);
        if (ret != RT_ERROR_NONE) {
            HCCL_ERROR("[Destroy][Event]destroy event failed.");
            return HCCL_E_RUNTIME;
        }
    }
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelInfoStore::GetDataTypeFromTaskInfo(const ge::GETaskInfo &task, HcclDataType &dataType) const
{
    HvdKernelInfoPrivateDef *privateDefBuf = static_cast<HvdKernelInfoPrivateDef *>(task.privateDef);
    if ((privateDefBuf->dataType >= HCCL_DATA_TYPE_INT8) &&
        (privateDefBuf->dataType < HCCL_DATA_TYPE_RESERVED)) {
        dataType = HcclDataType(privateDefBuf->dataType);
    } else {
        HCCL_ERROR("[Get][DataType]errNo[0x%016llx] get date type from task info failed. dataType[%s] is invalid.",
            HCOM_ERROR_CODE(HCCL_E_PARA), GetDataTypeEnumStr(privateDefBuf->dataType).c_str());
        return HCCL_E_PARA;
    }
    HCCL_INFO("get dataType[%s] from task info success. expect:[%d]-[%d]", GetDataTypeEnumStr(dataType).c_str(),
        HCCL_DATA_TYPE_INT8, HCCL_DATA_TYPE_RESERVED);
    return HCCL_SUCCESS;
}
}
